"""
Reinforcement Learning (A3C) using Pytroch + multiprocessing.
The most simple implementation for continuous action.

View more on my Chinese tutorial page [莫烦Python](https://morvanzhou.github.io/).
"""

from AI_pytorch_A3C.Tools import v_wrap, set_init, push_and_pull, record
import torch.multiprocessing as mp
from AI_pytorch_A3C.shared_adam import SharedAdam
import numpy as np
import math, os
from AI_pytorch_A3C.Network import Net
import Env.EnvMain as EM
import Env.Tools.MySQLFunc as MSF
os.environ["OMP_NUM_THREADS"] = "1"
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import axes3d
from AI_pytorch_A3C.IO_Tools import preprocess
import json
import time
import Env.Tools.MySQLFunc as MSF
import datetime

para=json.load(open('AI_pytorch_A3C/para.json','r',encoding='UTF-8'))
UPDATE_GLOBAL_ITER = para['全局更新迭代次数UPDATE_GLOBAL_ITER'] # k步更新
GAMMA = para['衰减GAMMA']
MAX_EP = para['最大回合数MAX_EP']
MAX_EP_STEP = para['回合最大步数MAX_EP_STEP']

Opponents=para['对手数量Opponents']
Friends=para['友方数量Friends']

N_S = para['状态维数N_S']     # 输入状态的维数
N_A = para['机动输出维数N_A']     # 输出的俯仰角、偏转角
N_AA=para['攻击输出维数N_AA']      # 输出的攻击动作，0为不攻击，1为打第一个目标


class Worker(mp.Process):
    def __init__(self, gnet, opt, global_ep, global_ep_r, res_queue, name,sqlname):
        super(Worker, self).__init__()
        self.name = 'w%i' % name
        self.g_ep, self.g_ep_r, self.res_queue = global_ep, global_ep_r, res_queue
        self.gnet, self.opt = gnet, opt
        self.lnet = Net(N_S, N_A,N_AA)           # local network
        self.env = EM.ACEnv(Opponents,Friends,index=name)
        self.sql_table_name=sqlname


    def run(self):
        total_step = 1
        sql=MSF.Mysql(self.name)

        while self.g_ep.value < MAX_EP:
            sql_collections = []
            with self.g_ep.get_lock():
                self.g_ep.value += 1
                ep = self.g_ep.value
            print('EP:',ep)
            self.env.reset()
            s=self.env.FirstStep(First_step_strategy())
            pre=preprocess(s)
            s=pre.get_input()
            buffer_s, buffer_a, buffer_r,buffer_aa = [], [], [],[]
            ep_r = 0.

            if self.name == 'w0':
                fig = plt.figure()
                self.env.move()
                fig.suptitle(self.name)
                ax = fig.add_subplot(111, projection='3d')

            for t in range(MAX_EP_STEP):
                if self.name == 'w0':
                     self.fig_show(fig,True)

                a = self.lnet.choose_action(v_wrap(s[None, :]))

                s_, r, done= self.env.step([[a[0].clip(-0.17, 0.17)],[a[1]]])
                pre = preprocess(s_)
                s_ = pre.get_input()
                if t == MAX_EP_STEP - 1:
                    done = True
                ep_r += r
                buffer_a.append(a[0])
                buffer_aa.append(a[1][0])
                buffer_s.append(s)
                buffer_r.append((r+8.1)/8.1)    # normalize



                if total_step % UPDATE_GLOBAL_ITER == 0 or not done:  # update global and assign to local net
                    # sync
                    push_and_pull(self.opt, self.lnet, self.gnet, done, s_, buffer_s, buffer_a,buffer_aa, buffer_r, GAMMA)
                    buffer_s, buffer_a, buffer_r,buffer_aa = [], [], [],[]

                    if not done:  # done and print information
                        record(self.g_ep, self.g_ep_r, ep_r, self.res_queue, self.name)
                        break
                s = s_
                total_step += 1
                decisions=[[float(a[0][0])],[float(a[0][1])],[float(a[1][0])],[float(0)],[float()]]
                sql_collections+=self.env.datacollect_mysql(ep,t,decisions)

            if self.name=='w0':
                plt.close()
            sql.Insert_record(self.sql_table_name,sql_collections)

        self.res_queue.put(None)    # 给出记录退出标记

    def fig_show(self,fig,show=False):
        if show:
            plt.clf()  # 清除前一回合图像
            ax = fig.add_subplot(111, projection='3d')  # 重新建立三维坐标
            self.env.render(fig, ax)

def First_step_strategy():      # FirstStep的策略
    strategy=[]
    select=np.random.randint(0,120)
    strategy.append([0,select])
    return strategy

def index_info(sqlname):
    start_time=datetime.datetime.now()
    n_friend=int(Friends)
    n_enemy=int(Opponents)
    max_ep=int(MAX_EP)
    max_step=int(MAX_EP_STEP)
    table_index=sql_name
    end_time=datetime.datetime.now()
    return (start_time,n_friend,n_enemy,max_ep,max_step,table_index,end_time)

if __name__ == "__main__":
    gnet = Net(N_S, N_A,N_AA)        # global network 构建全局网络
    gnet.share_memory()         # share the global parameters in multiprocessing 多进程中分享的全局参数
    opt = SharedAdam(gnet.parameters(), lr=0.0002)  # global optimizer  全局优化器
    global_ep, global_ep_r, res_queue = mp.Value('i', 0), mp.Value('d', 0.), mp.Queue()     # Value用来多进程共享参数，Queue用来交换对象
    Mysql=MSF.Mysql('Main')
    sql_name=Mysql.Creat_table()
    Mysql.Insert_index(index_info(sql_name))

    # parallel training 输入为共享的网络，共享优化器，全局共享参数*3，进程序号
    workers = [Worker(gnet, opt, global_ep, global_ep_r, res_queue, i,sql_name) for i in range(mp.cpu_count())]
    [w.start() for w in workers]
    res = []                    # record episode reward to plot
    while True:
        r = res_queue.get()
        if r is not None:
            res.append(r)
        else:
            break
    [w.join() for w in workers]

    Mysql.EndTime(datetime.datetime.now(),sql_name)
    Mysql.Close()

    import matplotlib.pyplot as plt
    plt.plot(res)
    plt.ylabel('Moving average ep reward')
    plt.xlabel('Step')
    plt.show()
    plt.savefig('Env/Records/figs/'+time.strftime("%m-%d-%H-%M", time.localtime())+'.eps')
